import numpy as np
import seaborn as sns
import matplotlib.pyplot as plt
from sklearn.metrics import confusion_matrix, classification_report

import numpy as np
import seaborn as sns
import matplotlib.pyplot as plt


def generate_confusion_matrix(val_acc, f1, class_dist, total_samples=3107):
    """
    根据验证准确率、F1分数和类别分布生成模拟混淆矩阵
    参数:
        val_acc: 验证准确率 (0.7789)
        f1: F1分数 (0.7769)
        class_dist: 类别比例列表 [0.6267, 0.1743, 0.1687, 0.0301]
        total_samples: 总样本数 (3107)
    """
    num_classes = len(class_dist)
    cm = np.zeros((num_classes, num_classes), dtype=int)

    # 计算每个类别的真实样本数
    class_counts = [int(total_samples * p) for p in class_dist]
    class_counts[-1] = total_samples - sum(class_counts[:-1])  # 处理舍入误差

    # 分配正确预测的样本 (基于验证准确率)
    correct_predictions = int(total_samples * val_acc)

    # 按F1分数分配各类别的正确预测数 (保持precision≈recall≈F1)
    for i in range(num_classes):
        cm[i, i] = int(class_counts[i] * f1)  # recall ≈ f1

    # 调整对角线使总正确数匹配
    current_correct = np.sum(np.diag(cm))
    adjustment = correct_predictions - current_correct
    if adjustment > 0:
        # 将多出的正确预测优先分配给多数类
        cm[0, 0] += int(adjustment * 0.6)
        cm[1, 1] += int(adjustment * 0.2)
        cm[2, 2] += int(adjustment * 0.15)
        cm[3, 3] += max(0, adjustment - int(adjustment * 0.95))
    else:
        # 减少正确预测数（按比例）
        for i in range(num_classes):
            reduction = int(abs(adjustment) * class_dist[i])
            cm[i, i] = max(0, cm[i, i] - reduction)

    # 分配错误预测的样本
    for i in range(num_classes):
        remaining_errors = class_counts[i] - cm[i, i]
        for j in range(num_classes):
            if i != j:
                # 错误按类别比例分配，多数类更容易被误判
                error_ratio = class_dist[j] / (1 - class_dist[i])
                cm[i, j] = int(remaining_errors * error_ratio)

        # 处理舍入误差
        cm[i, -1] = class_counts[i] - np.sum(cm[i, :])

    # 最终调整确保总和正确
    cm[-1, -1] = max(0, cm[-1, -1] + (total_samples - np.sum(cm)))

    return cm


# 使用您的参数生成矩阵
cm = generate_confusion_matrix(
    val_acc=0.8616,
    f1=0.8599,
    class_dist=[0.6267, 0.1743, 0.1687, 0.0301],
    total_samples=3107
)

# 验证样本总数
print(f"Total samples: {cm.sum()}")  # 应输出3107

# 高级可视化设置
plt.figure(figsize=(10, 8))
ax = sns.heatmap(cm, annot=True, fmt="d", cmap="Blues",
                linewidths=0.5, linecolor='grey',
                annot_kws={"size": 14, "color": "black"})

# 专业级格式设置
ax.set_xlabel('Predicted Label', fontsize=12, labelpad=12)
ax.set_ylabel('True Label', fontsize=12, labelpad=12)
ax.set_title('Confusion Matrix ',
            fontsize=14, pad=20, weight='bold')

# 设置刻度标签
class_names = ['Class 0', 'Class 1', 'Class 2', 'Class 3']
ax.set_xticklabels(class_names, rotation=45, ha='right')
ax.set_yticklabels(class_names, rotation=0)

# 添加关键指标标注


class_dist = cm.sum(axis=1)/cm.sum()


plt.tight_layout()
plt.show()

# 生成分类报告
print("\nDetailed Classification Report:")
print(classification_report(
    np.repeat([0,1,2,3], cm.sum(axis=1)),
    np.repeat([0,1,2,3], np.diag(cm)),
    target_names=class_names,
    digits=4
))